
import { TavilySearchResults } from "@langchain/community/tools/tavily_search";
import { ChatOpenAI } from "@langchain/openai";
import { HumanMessage, AIMessage } from "@langchain/core/messages";
import { ToolNode } from "@langchain/langgraph/prebuilt";
import { StateGraph, MessagesAnnotation } from "@langchain/langgraph";

// Define the tools for the agent to use
const tools = [new TavilySearchResults({ maxResults: 3 ,apiKey: "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"})];
const toolNode = new ToolNode(tools);


// const model = new ChatOpenAI({
//     model: 'chatgpt-4o',
//     configuration: {
//         apiKey: "sk-SrZfBqhIlUE8BBbb0d974049B0Fa4f29943419653176F06c",
//         baseURL: "https://pro.aiskt.com/v1",
//     },
//     temperature: 0,
//     verbose: true,
// });

// Create a model and give it access to the tools
const model  =  new ChatOpenAI({
    temperature: 0,
    model: "qwen2.5:7b",
    configuration: {
        apiKey: "test",
        verbose: true,
        baseURL: "http://127.0.0.1:11434/v1"
    }
}).bindTools(tools);

// Define the function that determines whether to continue or not
function shouldContinue({ messages }) {
    const lastMessage = messages[messages.length - 1];

    // If the LLM makes a tool call, then we route to the "tools" node
    if (lastMessage.tool_calls?.length) {
        return "tools";
    }
    // Otherwise, we stop (reply to the user) using the special "__end__" node
    return "__end__";
}

// Define the function that calls the model
async function callModel(state) {
    const response = await model.invoke(state.messages);

    // We return a list, because this will get added to the existing list
    return { messages: [response] };
}

// Define a new graph
const workflow = new StateGraph(MessagesAnnotation)
    .addNode("agent", callModel)
    .addEdge("__start__", "agent") // __start__ is a special name for the entrypoint
    .addNode("tools", toolNode)
    .addEdge("tools", "agent")
    .addConditionalEdges("agent", shouldContinue);

// Finally, we compile it into a LangChain Runnable.
const app = workflow.compile();

// Use the agent
const finalState = await app.invoke({
    messages: [new HumanMessage("what is the weather in sf")],
});
console.log(finalState.messages[finalState.messages.length - 1].content);

const nextState = await app.invoke({
    // Including the messages from the previous run gives the LLM context.
    // This way it knows we're asking about the weather in NY
    messages: [...finalState.messages, new HumanMessage("what about ny")],
});
console.log(nextState.messages[nextState.messages.length - 1].content);
